from typing import Union, Iterable, Optional

import aio_pika
from aio_pika.abc import AbstractIncomingMessage, AbstractQueue, ExchangeType, DateType
from aiormq.abc import ConfirmationFrameType


class RabbitMQUtil:

    def __init__(self, mq_cli, timeout=5):
        self.mq_cli = mq_cli
        if timeout is None:
            timeout = 5
        self.timeout = timeout

    async def declare_queue_and_bind_exchange(self,
                                              queue_name: str,
                                              exchange_name: str,
                                              routing_key: str,
                                              exchange_type: Union[ExchangeType, str] = ExchangeType.DIRECT,
                                              duration: bool = False
                                              ):
        async with self.mq_cli.acquire() as chan:
            queue = await chan.declare_queue(name=queue_name, durable=True)
            exchange = await self.declare_exchange(exchange_name, exchange_type, durable=duration)
            await queue.bind(exchange, routing_key=routing_key)

    async def declare_exchange(self, exchange_name, exchange_type, durable):
        async with self.mq_cli.acquire() as chan:
            exchange = await chan.declare_exchange(name=exchange_name, durable=durable, type=exchange_type)
            return exchange

    async def declare_queue(self, queue_name, duration, args=None):
        async with self.mq_cli.acquire() as chan:
            queue = await chan.declare_queue(name=queue_name, durable=duration, arguments=args)
            return queue

    async def send_msg(self,
                       msg: str,
                       routing_key: str,
                       exchange_name: str = None,
                       timeout: int = None
                       ):
        """
        发送消息到队列
        """
        async with self.mq_cli.acquire() as chan:  # type: aio_pika.abc.AbstractChannel
            exchange = chan.get_exchange(name=exchange_name) if exchange_name else chan.default_exchange
            msg = aio_pika.Message(body=msg.encode())
            timeout = timeout or self.timeout
            res = await exchange.publish(msg, routing_key=routing_key, timeout=timeout)
            return res

    async def get_msg(self, queue_name: str, timeout: int = 5) -> Union[AbstractIncomingMessage, None]:
        """
        从队列拉取消息
        """
        timeout = timeout or self.timeout
        async with self.mq_cli.acquire() as chan:
            queue = await chan.declare_queue(name=queue_name, durable=True, timeout=timeout)  # type: AbstractQueue
            msg = await queue.get(fail=False, timeout=timeout)
            return msg

    async def get_msgs(self, queue_name: str, qos: int = 10) -> Iterable:
        async with self.mq_cli.acquire() as chan:
            queue = await chan.declare_queue(name=queue_name, durable=True, timeout=self.timeout)  # type: AbstractQueue
            await chan.set_qos(qos)
            async with queue.iterator() as queue_iter:
                async for msg in queue_iter:
                    yield msg


class OrderMQ(RabbitMQUtil):

    queue_name = "order"
    exchange_name = "mall_order"
    routing_key = "order"
    exchange = None
    queue = None
    initd = False

    async def _init(self):
        if not self.initd:
            if not self.exchange:
                self.exchange = await self.declare_exchange(self.exchange_name, exchange_type=ExchangeType.DIRECT, durable=True)
            if not self.queue:
                self.queue = await self.declare_queue(self.queue_name, duration=True)
                await self.queue.bind(self.exchange, routing_key=self.routing_key)
            self.initd = True

    async def get_order_msg(self) -> AbstractIncomingMessage:
        await self._init()
        return await self.get_msg(self.queue_name)


class DelayMsgQueue(RabbitMQUtil):

    def __init__(self, mq_cli, exchange_name, routing_key, un_used_queue_name, dead_queue_name,
                 dead_exchange_name, dead_exchange_routing_key, ttl):
        super().__init__(mq_cli)
        self.exchange_name = exchange_name
        self.routing_key = routing_key
        self.un_used_queue_name = un_used_queue_name
        self.dead_exchange_name = dead_exchange_name
        self.dead_queue_name = dead_queue_name
        self.dead_exchange_routing_key = dead_exchange_routing_key
        self.ttl = ttl
        self._inited = False

    async def _init(self):
        """
        完成交换机和队列的初始化以及绑定
        Returns:
            None
        """
        if not self._inited:
            headers = {
                "x-dead-letter-exchange": self.dead_exchange_name,
                "x-dead-letter-routing-key": self.dead_exchange_routing_key
            }
            exchange = await self.declare_exchange(self.exchange_name, exchange_type=ExchangeType.DIRECT, durable=True)
            e = await self.declare_exchange(self.dead_exchange_name, exchange_type=ExchangeType.FANOUT, durable=True)
            q = await self.declare_queue(self.dead_queue_name, duration=True)
            await q.bind(e, routing_key=self.dead_exchange_routing_key)
            un_used_queue = await self.declare_queue(self.un_used_queue_name, duration=True, args=headers)
            await un_used_queue.bind(exchange, routing_key=self.routing_key)
            self._inited = True

    async def send_delay_msg(self,
                             msg: str,
                             timeout: Optional[int] = None
                             ) -> ConfirmationFrameType:
        await self._init()
        async with self.mq_cli.acquire() as chan:
            exchange = await chan.get_exchange(self.exchange_name)
            msg = aio_pika.Message(body=msg.encode("utf-8"), expiration=self.ttl)
            timeout = timeout or self.timeout
            return await exchange.publish(msg, routing_key=self.routing_key, timeout=timeout)





